#!/usr/bin/env python

import optparse
import os
import subprocess
import sys
import tempfile

from charmicat import BadCommandLineError

import ypipi

class Pi0Pi0SlurmScript(object):
    script_body = '''#!/bin/bash
#SBATCH -A belle
#SBATCH -t {duration}
#SBATCH -N {nnodes}
#SBATCH -o {stdout}
#SBATCH -e {stderr}

source $HOME/belle/belle_custom_setup.sh
source $HOME/belle/belle_general_setup.sh

cat {list_file} | xargs ypipi_pi0pi0_cuts.py -o {output} {cut_name}
'''

    def __init__(self, cut_name, list_file, outdir, duration, nnodes, job_name, stdoutdir,
                 stderrdir, stdout, stderr):
        
        output = os.path.join(os.path.abspath(outdir), cut_name + '.h5')
        
        list_file = os.path.abspath(list_file)
        stdoutdir = os.path.abspath(stdoutdir)
        stderrdir = os.path.abspath(stderrdir)
    
        if stdout is None:
            stdout = os.path.join(stderrdir, '{0}_%j.out'.format(job_name))
        
        if stderr is None:
            stderr = os.path.join(stderrdir, '{0}_%j.err'.format(job_name))

        self.script = Pi0Pi0SlurmScript.script_body.format(**locals())

    def __str__(self):
        return self.script

        


def main():
    '''
    ypipi_submit_cuts.py [opts] cut_set file.h5 [file.h5 [...]]

    Submit via `sbatch' jobs executing `ypipi_pi0pi0_cuts.py'; each job perform one 
    of the cuts in `cut_set' on all of the input files supplied.
    '''

    op = optparse.OptionParser(usage=main.__doc__)

    op.add_option("-l", "--list-file", dest='list_file', default='pi0pi0_cuts_file_list.txt',
                  help='Write input file list to PATH', metavar='PATH')
    op.add_option("-j", "--job-name", dest='job_name', default='pi0pi0_cuts',
                  help='Name jobs NAME in SLURM queue', metavar='NAME')
    op.add_option("-o", "--outdir", dest='outdir', default=os.getcwd(),
                  help='Leave output files in DIR', metavar='DIR')
    op.add_option("-t", "--duration", dest='duration', default='01:00:00',
                  help='Set job duration of N minutes or HH:MM:SS',
                  metavar='N (or) HH:MM:SS')
    op.add_option("--nodes", dest='nnodes', type='int', default=1,
                  help='Use N nodes per job', metavar='N')
    op.add_option("--stdout", dest='stdout', default=None,
                  help='Redirect job stdout to PATH', metavar='PATH')
    op.add_option("--stderr", dest='stderr', default=None,
                  help='Redirect job stderr to PATH', metavar='PATH')
    op.add_option("--stdout-dir", dest='stdoutdir', default=os.getcwd(),
                  help='Place job stdout in DIR', metavar='DIR')
    op.add_option("--stderr-dir", dest='stderrdir', default=os.getcwd(),
                  help='Place job stderr in DIR', metavar='DIR')

    opts, args = op.parse_args()

    if len(args) == 0:
        msg = "Please supply the name of a cutset and at least one input file."
        raise BadCommandLineError(msg, -1)

    cutset_name = args.pop(0)

    if len(args) == 0:
        msg = "Please supply at least one input file."
        raise BadCommandLineError(msg, -2)

    if cutset_name not in ypipi.cutsets:
        msg = "Cut set with name `{0}' not present in ypipi.cutsets".format(cutset_name)
        raise BadCommandLineError(msg, -3)

    cutset = ypipi.cutsets[cutset_name]

    with open(os.path.join(opts.outdir, opts.list_file), 'w') as fp:
        for path in args:
            fp.write(os.path.abspath(path) + '\n')


    for name in cutset:
        scr = Pi0Pi0SlurmScript(name, **vars(opts))
        
        with tempfile.NamedTemporaryFile() as fp:
#        with open(name + '.sh', 'w') as fp:
            fp.write(str(scr))
            fp.flush()

            sbatch_pipe = subprocess.Popen('sbatch -p shared'.split() + [fp.name], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
            
            stdout, stderr = sbatch_pipe.communicate()

            if sbatch_pipe.returncode != 0:
                msg = 'Error submitting job through sbatch:'
                msg += "\n\tSTDOUT: `{0}'".format(stdout)
                msg += "\n\tSTDERR: `{0}'".format(stderr)                
        
                print msg
                print "Continuing..."

                continue

            print stdout.strip()

    
    return 0



if __name__ == "__main__":
    try:
        sys.exit(main())
    except BadCommandLineError, e:
        import os
        sys.stderr.write(str(e) + "\n")
        sys.stderr.write("See `{0} -h' for more.\n".format(os.path.basename(sys.argv[0])))
        sys.stderr.flush()
        sys.exit(e.retcode)
